package com.xiashitech.agent.instrumentation.biz.desensitization.httpurl.filter.inst;

import com.xiashitech.agent.cache.status.CacheStatus;
import com.xiashitech.agent.config.StartConfig;
import com.xiashitech.agent.dto.UserRequest;
import com.xiashitech.agent.instrumentation.httprequest.XSHttpServletResponseWrapper;
import com.xiashitech.agent.instrumentation.utils.BaggageUtil;
import com.xiashitech.agent.utils.DurationPrinter;
import com.xiashitech.interfaces.agent.dto.AgentConfigDTO;
import io.opentelemetry.javaagent.bootstrap.CallDepth;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Scope;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletResponse;
import java.util.Map;

import static com.xiashitech.agent.instrumentation.biz.desensitization.httpurl.filter.XSServletSingletons.instrumenter;
import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext;
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
import static net.bytebuddy.matcher.ElementMatchers.*;

public class XSHttpFilterInstrumentation implements TypeInstrumentation {

    @Override
    public ElementMatcher<ClassLoader> classLoaderOptimization() {
        return hasClassesNamed("javax.servlet.Filter");
    }

    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {
        return hasSuperType(named("javax.servlet.Filter"));

    }

    @Override
    public void transform(TypeTransformer typeTransformer) {
        typeTransformer.applyAdviceToMethod(
                isMethod().and(named("doFilter"))
                        .and(takesArgument(0, named("javax.servlet.ServletRequest")))
                        .and(takesArgument(1, named("javax.servlet.ServletResponse")))
                        .and(takesArgument(2, named("javax.servlet.FilterChain")))
                        .and(isPublic()),
                XSHttpFilterInstrumentation.class.getName() + "$DesensitizationAdvice");
    }



    @SuppressWarnings("unused")
    public static class DesensitizationAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
                @Advice.Argument(value = 0, readOnly = false) ServletRequest request,
                @Advice.Argument(value = 1, readOnly = false) ServletResponse response,
                @Advice.Local("otelCallDepth") CallDepth callDepth,
                @Advice.Local("otelContext") Context context,
                @Advice.Local("otelScope") Scope scope,
                @Advice.Local("otelBaggageScope") Scope baggageScope,
                @Advice.Local("user") UserRequest userRequest,
                @Advice.Local("desenseImpl") String desenseImpl) {
            DurationPrinter durationPrinter = new DurationPrinter("XSHttpFilterInstrumentation");
            try {
                if (CacheStatus.cannotRun()) {
                    return;
                }
                Map<String, String> desenseImplMap = AgentConfigDTO.getConfig().getDesenseImpl();
                if (desenseImplMap != null){
                    desenseImpl = desenseImplMap.get(StartConfig.getParamValueFromKey(StartConfig.systemId));
                }
                if (desenseImpl == null || "".equals(desenseImpl) || !"servletFilter".equalsIgnoreCase(desenseImpl)){
                    return;
                }
                if (!(response instanceof XSHttpServletResponseWrapper)){
                    return;
                }
                userRequest = UserRequest.builder().build();
                userRequest.setRequest(request);
                userRequest.setAnyObject(response);
                Context parentContext = currentContext();
                if (instrumenter().shouldStart(parentContext, userRequest)) {
                    parentContext = BaggageUtil.syncBaggage().build().storeInContext(parentContext);
                    context = instrumenter().start(parentContext, userRequest);
                    context = BaggageUtil.syncBaggage().build().storeInContext(context);
                    if (response instanceof XSHttpServletResponseWrapper){
                        context = BaggageUtil.syncBaggage().put("endpoint-api", userRequest.getEndpointApi())
                                .put("api-systemId", StartConfig.getParamValueFromKey(StartConfig.systemId)).build()
                                .storeInContext(context);
                    }
                    scope = context.makeCurrent();
                }
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("DesensitizationAdvice b");
        }


        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
                                  @Advice.Argument(value = 0, readOnly = false) ServletRequest request,
                                  @Advice.Argument(value = 1, readOnly = false) ServletResponse response,
                                  @Advice.Local("otelCallDepth") CallDepth callDepth,
                                  @Advice.Local("otelContext") Context context,
                                  @Advice.Local("otelScope") Scope scope,
                                  @Advice.Local("user") UserRequest userRequest,
                                  @Advice.Local("desenseImpl") String desenseImpl) {
            DurationPrinter durationPrinter = new DurationPrinter("XSHttpFilterInstrumentation");
            try {
                if (CacheStatus.cannotRun()) {
                    return;
                }
                if (desenseImpl == null || "".equals(desenseImpl) || !"servletFilter".equalsIgnoreCase(desenseImpl)){
                    return;
                }
                if (!(response instanceof XSHttpServletResponseWrapper)){
                    return;
                }
                scope.close();
                instrumenter().end(context, userRequest, (HttpServletResponse) response, null);
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("DesensitizationAdvice e");
        }

    }

}
